#!/usr/bin/env python
"""
简单的处理脚本，专门用于从annotation/4.json文件中提取威胁情报数据。
"""

import json
import os
import logging
import sys
from dotenv import load_dotenv
sys.path.append('knowledge-graph/backend/src')
from integrations.annotation_helper import AnnotationProcessor
from integrations.threat_intelligence import ThreatIntelligenceExtractor

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# 加载环境变量
load_dotenv()

# Neo4j连接配置
NEO4J_URI = os.getenv("NEO4J_URI", "bolt://localhost:7687")
NEO4J_USER = os.getenv("NEO4J_USER", "neo4j")
NEO4J_PASSWORD = os.getenv("NEO4J_PASSWORD", "password")

def process_annotation_file(file_path='annotation/4.json'):
    """处理指定的标注文件并提取威胁情报数据"""
    try:
        # 初始化组件
        ti_extractor = ThreatIntelligenceExtractor(NEO4J_URI, NEO4J_USER, NEO4J_PASSWORD)
        annotation_processor = AnnotationProcessor()
        
        logger.info(f"处理文件: {file_path}")
        
        # 加载JSON数据
        with open(file_path, 'r', encoding='utf-8') as f:
            annotation_data = json.load(f)
        
        # 输出原始数据的基本信息以便调试
        logger.info(f"原始数据类型: {type(annotation_data)}")
        logger.info(f"原始数据键: {list(annotation_data.keys())}")
        
        # 预处理标注数据
        logger.info("开始预处理标注数据...")
        preprocessed_data = annotation_processor.preprocess_annotation(annotation_data)
        
        # 输出预处理数据信息
        logger.info(f"预处理数据包含 {len(preprocessed_data.get('entities', []))} 个实体和 "
                  f"{len(preprocessed_data.get('relationships', []))} 个关系")
        
        # 输出object_refs，这是STIX报告所必需的
        logger.info(f"预处理数据包含 {len(preprocessed_data.get('object_refs', []))} 个对象引用")
        
        # 将预处理数据保存到文件以便检查
        with open('preprocessed_data.json', 'w', encoding='utf-8') as f:
            json.dump(preprocessed_data, f, indent=2)
            logger.info("预处理数据已保存到 preprocessed_data.json")
        
        # 使用威胁情报提取器处理数据
        logger.info("使用威胁情报提取器处理数据...")
        result = ti_extractor.process_json_data(preprocessed_data)
        
        # 输出结果
        if result.get("success", False):
            logger.info(f"成功提取了 {result.get('entity_count', 0)} 个实体和 "
                      f"{result.get('relation_count', 0)} 个关系")
            
            # 保存结果到JSON文件以便检查
            output_file = "annotation_result.json"
            with open(output_file, 'w', encoding='utf-8') as f:
                json.dump(result, f, indent=2)
            
            logger.info(f"结果已保存到 {output_file}")
            return True
        else:
            logger.error(f"提取实体失败: {result.get('error', '未知错误')}")
            return False
    
    except Exception as e:
        logger.error(f"处理标注文件时出错: {str(e)}")
        import traceback
        logger.error(traceback.format_exc())
        return False

if __name__ == "__main__":
    # 如果指定了文件参数则使用该文件，否则默认使用annotation/4.json
    file_path = sys.argv[1] if len(sys.argv) > 1 else 'annotation/4.json'
    success = process_annotation_file(file_path)
    if success:
        print("处理成功！")
    else:
        print("处理失败！")
        sys.exit(1) 